/////////////////////////////////////////////////////////////////
// SparseMatrix.h
//
// Sparse matrix computations
/////////////////////////////////////////////////////////////////

#ifndef SPARSEMATRIX_H
#define SPARSEMATRIX_H

#include <iostream>

using namespace std;

namespace PROBCONS {

  const float POSTERIOR_CUTOFF = 0.01;         // minimum posterior probability
  // value that is maintained in the
  // sparse matrix representation

  typedef pair<int,float> PIF;                 // Sparse matrix entry type
  //   first --> column
  //   second --> value

  /////////////////////////////////////////////////////////////////
  // SparseMatrix
  //
  // Class for sparse matrix computations
  /////////////////////////////////////////////////////////////////

  class SparseMatrix {

    int seq1Length, seq2Length;                     // dimensions of matrix
    VI rowSize;                                     // rowSize[i] = # of cells in row i
    SafeVector<PIF> data;                           // data values
    SafeVector<SafeVector<PIF>::iterator> rowPtrs;  // pointers to the beginning of each row

    /////////////////////////////////////////////////////////////////
    // SparseMatrix::SparseMatrix()
    //
    // Private constructor.
    /////////////////////////////////////////////////////////////////

    SparseMatrix (){}

  public:

    /////////////////////////////////////////////////////////////////
    // SparseMatrix::SparseMatrix()
    //
    // Constructor.  Builds a sparse matrix from a posterior matrix.
    // Note that the expected format for the posterior matrix is as
    // a (seq1Length+1) x (seq2Length+1) matrix where the 0th row
    // and 0th column are ignored (they should contain all zeroes).
    /////////////////////////////////////////////////////////////////

    SparseMatrix (int seq1Length, int seq2Length, const VF &posterior) :
      seq1Length (seq1Length), seq2Length (seq2Length) {

      int numCells = 0;

      assert (seq1Length > 0);
      assert (seq2Length > 0);

      // calculate memory required; count the number of cells in the
      // posterior matrix above the threshold
      VF::const_iterator postPtr = posterior.begin();
      for (int i = 0; i <= seq1Length; i++){
	for (int j = 0; j <= seq2Length; j++){
	  if (*(postPtr++) >= POSTERIOR_CUTOFF){
	    assert (i != 0 && j != 0);
	    numCells++;
	  }
	}
      }
    
      // allocate memory
      data.resize(numCells);
      rowSize.resize (seq1Length + 1); rowSize[0] = -1;
      rowPtrs.resize (seq1Length + 1); rowPtrs[0] = data.end();

      // build sparse matrix
      postPtr = posterior.begin() + seq2Length + 1;           // note that we're skipping the first row here
      SafeVector<PIF>::iterator dataPtr = data.begin();
      for (int i = 1; i <= seq1Length; i++){
	postPtr++;                                            // and skipping the first column of each row
	rowPtrs[i] = dataPtr;
	for (int j = 1; j <= seq2Length; j++){
	  if (*postPtr >= POSTERIOR_CUTOFF){
	    dataPtr->first = j;
	    dataPtr->second = *postPtr;
	    dataPtr++;
	  }
	  postPtr++;
	}
	rowSize[i] = dataPtr - rowPtrs[i];
      }
    }

    /////////////////////////////////////////////////////////////////
    // SparseMatrix::GetRowPtr()
    //
    // Returns the pointer to a particular row in the sparse matrix.
    /////////////////////////////////////////////////////////////////

    SafeVector<PIF>::iterator GetRowPtr (int row) const {
      assert (row >= 1 && row <= seq1Length);
      return rowPtrs[row];
    }

    /////////////////////////////////////////////////////////////////
    // SparseMatrix::GetValue()
    //
    // Returns value at a particular row, column.
    /////////////////////////////////////////////////////////////////

    float GetValue (int row, int col){
      assert (row >= 1 && row <= seq1Length);
      assert (col >= 1 && col <= seq2Length);
      for (int i = 0; i < rowSize[row]; i++){
	if (rowPtrs[row][i].first == col) return rowPtrs[row][i].second;
      }
      return 0;
    }

    /////////////////////////////////////////////////////////////////
    // SparseMatrix::GetRowSize()
    //
    // Returns the number of entries in a particular row.
    /////////////////////////////////////////////////////////////////

    int GetRowSize (int row) const {
      assert (row >= 1 && row <= seq1Length);
      return rowSize[row];
    }

    /////////////////////////////////////////////////////////////////
    // SparseMatrix::GetSeq1Length()
    //
    // Returns the first dimension of the matrix.
    /////////////////////////////////////////////////////////////////

    int GetSeq1Length () const {
      return seq1Length;
    }

    /////////////////////////////////////////////////////////////////
    // SparseMatrix::GetSeq2Length()
    //
    // Returns the second dimension of the matrix.
    /////////////////////////////////////////////////////////////////

    int GetSeq2Length () const {
      return seq2Length;
    }

    /////////////////////////////////////////////////////////////////
    // SparseMatrix::GetRowPtr
    //
    // Returns the pointer to a particular row in the sparse matrix.
    /////////////////////////////////////////////////////////////////

    int GetNumCells () const {
      return data.size();
    }

    /////////////////////////////////////////////////////////////////
    // SparseMatrix::Print()
    //
    // Prints out a sparse matrix.
    /////////////////////////////////////////////////////////////////

    void Print (ostream &outfile) const {
      outfile << "Sparse Matrix:" << endl;
      for (int i = 1; i <= seq1Length; i++){
	outfile << "  " << i << ":";
	for (int j = 0; j < rowSize[i]; j++){
	  outfile << " (" << rowPtrs[i][j].first << "," << rowPtrs[i][j].second << ")";
	}
	outfile << endl;
      }
    }

    /////////////////////////////////////////////////////////////////
    // SparseMatrix::ComputeTranspose()
    //
    // Returns a new sparse matrix containing the transpose of the
    // current matrix.
    /////////////////////////////////////////////////////////////////

    SparseMatrix *ComputeTranspose () const {

      // create a new sparse matrix
      SparseMatrix *ret = new SparseMatrix();
      int numCells = data.size();

      ret->seq1Length = seq2Length;
      ret->seq2Length = seq1Length;

      // allocate memory
      ret->data.resize (numCells);
      ret->rowSize.resize (seq2Length + 1); ret->rowSize[0] = -1;
      ret->rowPtrs.resize (seq2Length + 1); ret->rowPtrs[0] = ret->data.end();

      // compute row sizes
      for (int i = 1; i <= seq2Length; i++) ret->rowSize[i] = 0;
      for (int i = 0; i < numCells; i++)
	ret->rowSize[data[i].first]++;

      // compute row ptrs
      for (int i = 1; i <= seq2Length; i++){
	ret->rowPtrs[i] = (i == 1) ? ret->data.begin() : ret->rowPtrs[i-1] + ret->rowSize[i-1];
      }

      // now fill in data
      SafeVector<SafeVector<PIF>::iterator> currPtrs = ret->rowPtrs;

      for (int i = 1; i <= seq1Length; i++){
	SafeVector<PIF>::iterator row = rowPtrs[i];
	for (int j = 0; j < rowSize[i]; j++){
	  currPtrs[row[j].first]->first = i;
	  currPtrs[row[j].first]->second = row[j].second;
	  currPtrs[row[j].first]++;
	}
      }

      return ret;
    }

    /////////////////////////////////////////////////////////////////
    // SparseMatrix::GetPosterior()
    //
    // Return the posterior representation of the sparse matrix.
    /////////////////////////////////////////////////////////////////

    VF *GetPosterior () const {

      // create a new posterior matrix
      VF *posteriorPtr = new VF((seq1Length+1) * (seq2Length+1)); assert (posteriorPtr);
      VF &posterior = *posteriorPtr;

      // build the posterior matrix
      for (int i = 0; i < (seq1Length+1) * (seq2Length+1); i++) posterior[i] = 0;
      for (int i = 1; i <= seq1Length; i++){
	VF::iterator postPtr = posterior.begin() + i * (seq2Length+1);
	for (int j = 0; j < rowSize[i]; j++){
	  postPtr[rowPtrs[i][j].first] = rowPtrs[i][j].second;
	}
      }

      return posteriorPtr;
    }

  };
}
#endif
